{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import  tensorflow as tf\n",
    "from    tensorflow.keras import datasets, layers, optimizers, Sequential, metrics\n",
    "\n",
    "\n",
    "def preprocess(x, y):\n",
    "    \"\"\"\n",
    "    x is a simple image, not a batch\n",
    "    \"\"\"\n",
    "    x = tf.cast(x, dtype=tf.float32) / 255.\n",
    "    x = tf.reshape(x, [28*28])\n",
    "    y = tf.cast(y, dtype=tf.int32)\n",
    "    y = tf.one_hot(y, depth=10)\n",
    "    return x,y\n",
    "\n",
    "\n",
    "batchsz = 128\n",
    "(x, y), (x_val, y_val) = datasets.mnist.load_data()\n",
    "print('datasets:', x.shape, y.shape, x.min(), x.max())\n",
    "\n",
    "\n",
    "\n",
    "db = tf.data.Dataset.from_tensor_slices((x,y))\n",
    "db = db.map(preprocess).shuffle(60000).batch(batchsz)\n",
    "ds_val = tf.data.Dataset.from_tensor_slices((x_val, y_val))\n",
    "ds_val = ds_val.map(preprocess).batch(batchsz) \n",
    "\n",
    "sample = next(iter(db))\n",
    "print(sample[0].shape, sample[1].shape)\n",
    "\n",
    "\n",
    "network = Sequential([layers.Dense(256, activation='relu'),\n",
    "                     layers.Dense(128, activation='relu'),\n",
    "                     layers.Dense(64, activation='relu'),\n",
    "                     layers.Dense(32, activation='relu'),\n",
    "                     layers.Dense(10)])\n",
    "network.build(input_shape=(None, 28*28))\n",
    "network.summary()\n",
    "\n",
    "#compile\n",
    "#fit\n",
    "#evaluate\n",
    "#predict\n",
    "\n",
    "\n",
    "network.compile(optimizer=optimizers.Adam(lr=0.01),\n",
    "\t\tloss=tf.losses.CategoricalCrossentropy(from_logits=True),\n",
    "\t\tmetrics=['accuracy']\n",
    "\t)\n",
    "\n",
    "network.fit(db, epochs=5, validation_data=ds_val, validation_freq=2)\n",
    " \n",
    "network.evaluate(ds_val)\n",
    "\n",
    "sample = next(iter(ds_val))\n",
    "x = sample[0]\n",
    "y = sample[1] # one-hot\n",
    "pred = network.predict(x) # [b, 10]\n",
    "# convert back to number \n",
    "y = tf.argmax(y, axis=1)\n",
    "pred = tf.argmax(pred, axis=1)\n",
    "print('x.shape',x.shape)\n",
    "print('pred',pred)\n",
    "print('y',y)\n",
    "\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
